/*
The main file for DDMCMC algorithm.
*/

#include "ddmcmc.h"

DDMCMC *DDMCMCActive;


DDMCMC::DDMCMC() {
	/* Initialize a DDMCMC object. 
	 * Here we assume that the image is a color image, hence the scale(in ddmcmcSettings) is set to 3.0 by default.
	 */
	msClustering = NULL;
	pmap = NULL;
	m_pmcmc_conductor = NULL;
}

DDMCMC::~DDMCMC() {
}

void DDMCMC::generateProposal(Raster_Color<uchar> &image_color, bool visualize) {
	Mat					mPartitionMaps[3];
	mcmcMatrix<int>		mx_region_labels;
	CEdgePrcs			*pedge_linking_proc = NULL;

	double sig[3] = { 6.0, 3.0, 2.0 };

	image_color.CreateLUVspace();

	for (int k = 0; k < 3; k++) {
		pedge_linking_proc = new CEdgePrcs();
		num[k] = pedge_linking_proc->LabelRegions(sig[k], mx_region_labels, 1, &image_color);

		mPartitionMaps[k] = mx_region_labels.data;
		//mPartitionMaps[k].setTo(Scalar::all(1)); // Jiadi 
		//mPartitionMaps[k].at<uchar>(0, 0) = 0;
		delete pedge_linking_proc;

		partition_maps[k].SetDimension(mPartitionMaps[k].rows, mPartitionMaps[k].cols);
		partition_maps[k].data = mPartitionMaps[k].clone();
	}

	// Mean shift
	msClustering = new DDMCMC_MeanShiftClustering(mImg);
	msClustering->meanShiftSegmentation(1);

	Mat labels = msClustering->getLabels();
	ms_clustering.SetDimension(labels.rows, labels.cols);
	labels.convertTo(ms_clustering.data, CV_32S);

	if (visualize){
		Mat clusters = msClustering->getClusteredImage();
		cvtColor(clusters, clusters, CV_RGB2BGR);
		imshow("Mean Shift Clustering", clusters);
		waitKey(1);
	}

	return;
}

Mat DDMCMC::synthesize(){
	Raster_Color<uchar> raster_syn;
	Mat mSyn;
	pmap->Synthesize(raster_syn, 1);
	mSyn = raster_syn.RGB.clone();
	cvtColor(mSyn, mSyn, CV_RGB2BGR);
	return mSyn;
}

void DDMCMC::RunDDMCMCSegmentation(Raster_Color<uchar> &raster, mcmcMatrix<int> imageLabel[3], int num_label[], Raster<int> &image_init, const MCMC_SegSetting& setting1, char *str_file_name){
	Raster<int> matrix_temp;
	DDMCMCActive = this;

	// set seed for random value generator
	ztu_srand(time(NULL));

	// copy label image
	matrix_temp = image_init;

	// initialization
	// Siyuan: !! In the original program, pmap is recoverable
	// If it is not null: pmap->LayerRegion()->Recover(matrix_temp);
	pmap = new MCMC_Map_2Layers();

	matrix_temp.RelabelReorder();
	matrix_temp.RelabelAndMergeTinyRegions(5);

	try{
		pmap->LayerRegion()->Build(matrix_temp);
	}catch (XPT_OutOfBound *pxpt){
		delete pxpt;
		assert(false);
	}
	pmap->LayerRegion()->MergeSmallAtomicRegion();
	pmap->LayerRegion()->CreateTopologyRecords();
	// create group regions
	pmap->LayerRegion()->PutAllAtomicRegionsIntoGRegions();

	m_pmcmc_conductor = new MCMC_Conductor(*pmap, imageLabel, num_label, raster, ms_clustering, setting1, str_file_name);
	tStart = GetTickCount();

	// segmenting
	m_pmcmc_conductor->GoNew(str_file_name, setting1);

	mx_labelmap.SetDimension(raster.rows(), raster.cols());
	pmap->LayerRegion()->LabelAtomicRegions(mx_labelmap);

	// show the gradient image
	Raster_Color<uchar> raster_syn;
	pmap->Synthesize(raster_syn, 1);
	mSynResult = raster_syn.RGB.clone();
	cvtColor(mSynResult, mSynResult, CV_RGB2BGR);

	//show the final segmentation which consists of a lot of GRegions
	Raster_Color<uchar> ras_copy;
	Raster<int>			mx_boundary;
	ras_copy = raster;
	mx_boundary.SetDimension(raster.rows(), raster.cols());
	mx_boundary.InitValue(0);
	pmap->LayerRegion()->DrawRegions(mx_boundary);
	mBoundaries = Mat(raster.rows(), raster.cols(), CV_8U, Scalar(0));
	for (int j = 0; j < raster.rows(); j++){
		for (int i = 0; i<raster.cols(); i++){
			if (mx_boundary(j, i)>0){
				ras_copy.R(j, i) = 255;
				ras_copy.G(j, i) = 0;
				ras_copy.B(j, i) = 0;
				mBoundaries.at<uchar>(j, i) = 255;
			}
		}
	}
	mResult = ras_copy.RGB.clone();
	cvtColor(mResult, mResult, CV_RGB2BGR);
}

void DDMCMC::clearWorkspace(){
	mImg.release();
	mResult.release();
	mSynResult.release();
	mBoundaries.release();

	delete msClustering;
	delete pmap;
	delete m_pmcmc_conductor;

	msClustering = NULL;
	pmap = NULL;
	m_pmcmc_conductor = NULL;
}

void DDMCMC::settingInitialize(int stepFrom, int stepTo, double scale, bool visualize){
	setting.bcolor = true;
	setting.nstep_from = stepFrom;
	setting.nstep_to = stepTo;
	setting.dscale = scale;
	setting.visualize = visualize;
}

void DDMCMC::setImg(Mat mInputimg) {
	cvtColor(mInputimg, mImg, CV_BGR2RGB);
	if (mImg.cols*mImg.rows >= PIXELS_1080P) {
		resize(mImg, mImg, Size(), 0.5, 0.5);
	}
	else if (mImg.cols*mImg.rows >= PIXELS_720P) {
		resize(mImg, mImg, Size(), 0.75, 0.75);
	}
}

bool DDMCMC::DDMCMCSegmentation(Mat mInputimg, int stepFrom, int stepTo, double scale, bool visualize) {
	if (!mInputimg.data){
		cout << "Invalid input image!" << endl;
		return false;
	}
	
	// Initialize
	clearWorkspace();
	settingInitialize(stepFrom, stepTo, scale, visualize);
	setImg(mInputimg);

	Raster_Color<uchar> image_color;
	image_color.SetDimension(mImg.rows, mImg.cols);
	image_color.RGB = mImg.clone();

	// Generate proposal
	generateProposal(image_color, visualize);
	tStart = GetTickCount();

	Raster<int> image_init;
	image_init.SetDimension(ms_clustering.rows(), ms_clustering.cols());
	image_init = partition_maps[1];

	// Run DDMCMC segmentation
	RunDDMCMCSegmentation(image_color, partition_maps, num, image_init, setting);

	return true;
}

DDMCMC_Settings::DDMCMC_Settings() {
	bTexture = false;
	bSaveSeg = true;
	dScale = 3.0;
	nStepFrom = 0;
	nStepTo = 10;
	bSameSize = true;
	bLoadData = false;
}